#include "TimeSpaceMRF.h"

static float getSimilarity( Mat & desc1, Mat & desc2)
{
	float ans = 1.0f;

	float smooth = 0.7f;

	float max_diff = 0.25f;

	for(int i = 2;i <5 ; i++)
	{
		float diff = fabs(desc1.at<float>( 0,i) - desc2.at<float>( 0,i));				

		if(diff>max_diff) ans *= (smooth/(1+smooth));
		else
		{
			diff = max_diff - diff;
			ans *= (( diff  / max_diff + smooth)/(1+smooth));//* diff /max_diff
		}
	}	

	return ans;
}

void TimeSpaceMRF::saveRess( vector< Mat > & ress )
{
	for(int i = 0 ; i <(int) ress.size(); i++)
	{
		Mat dsp = ress[i] * 255.0;

		dsp.convertTo( dsp, CV_8U);
		char fname[400];
		sprintf( fname ,"./visual/res_%d.jpg", i);
		imwrite( fname , dsp );				
	}
}



void TimeSpaceMRF::buildAndCut(		vector< Mat > & ress,
									vector< int > & n_nodes,
									vector< int > & node_shift,
									vector< Mat >  & node_descs,									
									vector< Mat >  & label_maps,
									vector< SuperpixelFlow> & superpixel_flows,
									vector< Mat >  & adj_mats
									)
{	

	vector< Mat > node_weights;

	getNodeWeight( node_descs, node_weights);	

	vector< Mat > node_ress;

	node_ress.resize( n_nodes.size() );
		 ress.resize( n_nodes.size() );

	int n_all_node = 0;

	for(int i = 0 ; i < (int)  n_nodes.size() ; i++) 
		n_all_node += n_nodes[i];

	int n_all_edge = 0;

	for(int i =0; i < (int) superpixel_flows.size() ; i++)
		n_all_edge += int(superpixel_flows[i].second.size());

	for(int i = 0 ; i < (int) adj_mats.size(); i++)
	{
		Mat temp;
		adj_mats[i].convertTo( temp, 5);
		n_all_edge += int(lc::reduceMat( temp, CV_REDUCE_SUM));
	}

	float MAX_EDGE_CAP = 1000;

	typedef Graph< float , float , float > GraphType;

	GraphType * _g = new GraphType( n_all_node, n_all_edge * 2 );

	for(int i = 0 ; i < (int)node_weights.size(); i++)
	{
		int id = node_shift[i];
		for(int j = 0 ; j <node_weights[i].rows ; j++, id++)
		{
			_g->add_node();
			_g->add_tweights( id , node_weights[i].at<float>(j,0) + MAX_EDGE_CAP, MAX_EDGE_CAP );
		}
	}

	if(lambda > 1e-4f)
	{
		for(int i = 0 ; i < (int)adj_mats.size() ; i++)
		{
			int _n = adj_mats[i].rows;

			for(int x = 0; x < _n-1; x++)
				for(int y = x+1; y< _n ;y++)
					if(adj_mats[i].at<int>(x,y)!=0)
					{
						int id1 = node_shift[i] + x;
						int id2 = node_shift[i] + y;

						Mat desc1 = node_descs[i].row( x);
						Mat desc2 = node_descs[i].row( y);
						float val = getSimilarity(desc1 ,desc2);

						val *= lambda;

						_g->add_edge( id1, id2, val, val );

						_g->add_tweights( id1 , eta*val ,0  );
						_g->add_tweights( id2 , eta*val ,0  );
					}
		}
	}

	if(nu > 1e-4f )
	{
		for(int i =0; i < (int) superpixel_flows.size() ; i++)
		{
			int frame_id1 = superpixel_flows[i].first.x;
			int frame_id2 = superpixel_flows[i].first.y;

			vector<Point2i> & matches = superpixel_flows[i].second;

			for(int j = 0 ; j < (int) matches.size() ; j++)
			{
				int id1 = node_shift[ frame_id1 ] + matches[j].x;
				int id2 = node_shift[ frame_id2 ] + matches[j].y;

				Mat desc1 = node_descs[frame_id1].row( matches[j].x);
				Mat desc2 = node_descs[frame_id2].row( matches[j].y);
				float val = getSimilarity( desc1,desc2);

				val *= nu;

				_g->add_edge( id1, id2, val, val );

				_g->add_tweights( id1 , eta*val ,0  );
				_g->add_tweights( id2 , eta*val ,0  );
			}
		}

	}

	float flow = _g->maxflow();

	for(int i = 0; i < (int) node_ress.size();i++)
	{
		node_ress[i] = Mat::zeros( n_nodes[i] , 1, 5);

		int id = node_shift[i];

		for(int j = 0 ; j <n_nodes[i] ; j++, id++)
		{
			if( _g->what_segment(id) == GraphType::SOURCE) node_ress[i].at<float>(j, 0)= 1.0f;
			else node_ress[i].at<float>(j, 0)= 0.0f;
		}
		superResReconstruct(node_ress[i], label_maps[i], ress[i]);
	}

	delete _g;
}

TimeSpaceMRF::TimeSpaceMRF()
{
	IF_VISUAL = false;
	lambda = 0.04f;
	nu = 0.04f;
	heat_map_weight = 0.04f;
}

void TimeSpaceMRF::getAllInfo(	vector< Mat > & frms,
								LocalSensitiveBagging & crank_bagging,
								LcHistogram & global_extractor,
								vector< int > & n_nodes,
								vector< int > & node_shift,
								vector< Mat >  & node_descs,						
								vector< Mat >  & label_maps,
								vector< SuperpixelFlow> & superpixel_flows,
								vector< Mat >  & adj_mats)
{
	n_frame = (int) frms.size();
	t_win = n_frame /2;

	vector<Mat> likeli;
	vector< vector<Point2f> > node_centers;

	getLikeli( frms, likeli , crank_bagging, global_extractor);

	getSLIC( frms, likeli, label_maps, node_descs, adj_mats,
			 n_nodes, node_shift, node_centers );

	

	getSuperpixelFlow( superpixel_flows, frms , label_maps,node_descs, node_centers);

	if(IF_VISUAL) saveLikeli( likeli );

	if(IF_VISUAL) saveSuperpixelFlow( n_nodes, node_shift, node_centers, superpixel_flows);
}

LcValidator TimeSpaceMRF::predict(	vector< int > & n_nodes,
									vector< int > & node_shift,
									vector< Mat >  & node_descs,						
									vector< Mat >  & label_maps,
									vector< SuperpixelFlow> & superpixel_flows,
									vector< Mat >  & adj_mats, 
									Mat & gt,
									vector< Mat > & ress)
{
	buildAndCut( ress, n_nodes, node_shift, node_descs, label_maps,
		superpixel_flows, adj_mats );

	n_frame = (int) label_maps.size();
	t_win = n_frame /2;

	Mat res = ress[ t_win ];

	if( res.rows == gt.rows)
	{
		Mat _gt;
		if(gt.type()==CV_8U)
		{
			gt.convertTo( _gt , 5);
			_gt = _gt * (1.0/255.0);
		}
		else _gt = gt;

		return LcValidator( res, _gt);
	}
	else
		return LcValidator(0,0,0,0);
	
}

void TimeSpaceMRF::work(	vector< Mat > & frms,
							vector< Mat > & ress,				  
							LocalSensitiveBagging & crank_bagging,
							LcHistogram & global_extractor)
{	
	if(IF_VISUAL) saveSrc( frms );

	vector< Mat > label_maps;
	vector< Mat > node_descs;
	vector< Mat > adj_mats;
	vector< int > n_nodes;
	vector< int > node_shift;
	vector< SuperpixelFlow> superpixel_flows;

	getAllInfo( frms, crank_bagging, global_extractor, n_nodes,
				node_shift, node_descs, label_maps,
				superpixel_flows, adj_mats );

	Mat empty_gt = Mat::Mat();

	predict(	n_nodes,
				node_shift, node_descs, label_maps,
				superpixel_flows, adj_mats,
				empty_gt, ress );

	if(IF_VISUAL) saveRess( ress );

}

void TimeSpaceMRF::getNodeWeight(	vector< Mat >  & node_descs,
									vector< Mat > & node_weights)
{	
	//node_descs to node_weights
	{
		int _frame_n = (int) node_descs.size();

		node_weights.resize( _frame_n );

		for(int i = 0; i < _frame_n ; i++)
		{
			if(heat_map_weight<1e-4f)
				node_descs[i].col(0).copyTo( node_weights[i] );
			else
				node_weights[i] = node_descs[i].col(0) + node_descs[i].col(1) * double(heat_map_weight);
		}
	}
}

void TimeSpaceMRF::saveSuperpixelFlow(	vector<int> & n_nodes,
										vector< int > & node_shift,
										vector< vector<Point2f> > & node_centers,
										vector< SuperpixelFlow> & superpixel_flows)
{
	int n_all_node = 0;

	for(int i = 0 ; i < n_nodes.size() ; i++) 
		n_all_node += n_nodes[i];

	Mat node_save = Mat::zeros( n_all_node, 3, 5 );

	for(int i = 0 ; i < (int)n_nodes.size(); i++)
	{
		int id = node_shift[i];
		for(int j = 0 ; j <(int) node_centers[i].size() ; j++, id++)
		{
			node_save.at< float>(id,0) = float(i);
			node_save.at< float>(id,1) = node_centers[i][j].x;
			node_save.at< float>(id,2) = node_centers[i][j].y;
		}
	}

	lc::LcMat2Bin( "./visual/node_save.bin", node_save );

	int n_all_edge = 0;
	for(int i =0; i < superpixel_flows.size() ; i++)
		n_all_edge += int(superpixel_flows[i].second.size());

	Mat edge_save = Mat::zeros( n_all_edge, 2 , 5);

	for(int i =0,k=0; i < int(superpixel_flows.size()) ; i++)
	{
		int id1 = superpixel_flows[i].first.x;
		id1 = node_shift[id1];
		int id2 = superpixel_flows[i].first.y;
		id2 = node_shift[id2];

		for(int j = 0 ; j< int(superpixel_flows[i].second.size()); j++,k++)
		{
			edge_save.at<float>(k ,0) = float(id1 + superpixel_flows[i].second[j].x);
			edge_save.at<float>(k ,1) = float(id2 + superpixel_flows[i].second[j].y);
		}
	}

	lc::LcMat2Bin( "./visual/edge_save.bin", edge_save );
	
}


Match TimeSpaceMRF::getMatch(	int id1 , int id2 ,
								vector<Mat > & frms,
								vector< vector< Point2f > > & key_points,
								LcKLT & my_tracker)
{

	int _n1 = (int)key_points[id1].size();
	int _n2 = (int)key_points[id2].size();

	Match new_matches( Point2i( id1, id2), PointMaps());

	vector< Point2f > po1;
	vector< Point2f > po2;
	vector< int > sel_id;
	vector< int > out_id;

	if(_n1 > _n2 )
	{
		new_matches.second.first = key_points[id1];
		my_tracker.track( frms[id1], frms[id2],
						   new_matches.second.first,new_matches.second.second,
						   sel_id, out_id);
	}
	else
	{
		new_matches.second.second = key_points[id2];
		my_tracker.track( frms[id2], frms[id1],
						  new_matches.second.second,new_matches.second.first,
						  sel_id, out_id);
	}

	{
		vector<Point2f> & po1 = new_matches.second.first;
		vector<Point2f> & po2 = new_matches.second.second;

		int _n = int( sel_id.size());

		for(int j = 0 ; j < _n; j++)
		{
			po1[j] = po1[ sel_id[j] ];
			po2[j] = po2[ sel_id[j] ];
		}

		po1.resize( _n );
		po2.resize( _n );
	}

	return new_matches;
}


void TimeSpaceMRF::getMatches(vector< Mat > & frms, vector<Match> & matches)
{
	LcKLT my_tracker;

	my_tracker.veb = 0;

	vector< vector< Point2f > > key_points;

	getGoodPoints( frms, key_points, my_tracker);

	vector<int> valid_interval; valid_interval.clear();
	valid_interval.push_back(1); valid_interval.push_back(2);

	matches.reserve( n_frame * (int) valid_interval.size() );

	int count_tracked = 0;

	//========= do matching  ====================

	for(int k = 0 ; k < (int) valid_interval.size() ; k++)
	{
		int _d = valid_interval[k];

		for(int i = _d; i < n_frame; i++)
		{			

			Match new_matches = getMatch( i-_d, i , frms, key_points,my_tracker );

			count_tracked += int( new_matches.second.first.size() );

			matches.push_back( new_matches );

		}// end for frame

	}//end for valid_interval

	cout << count_tracked << " keypoint-pair success to track " << endl;	

}

void TimeSpaceMRF::getMeanOpt( vector< Point2f > & mean_opt1,Mat & label_map1, vector< Point2f > & po1,vector< Point2f > & po2, int n1,float frame_diff)
{
	float frame_speed_limit = 70.0f;

	vector< float > count_kp( n1, 1e-3f);

	int _rows = label_map1.rows;
	int _cols = label_map1.cols;

	//========= count optical flow in superpixel  =================			

	for(int i = 0 ; i < po1.size(); i++)
	{// sum over				
		int _y = int(po1[i].y); if(_y < 0 ) _y = 0; if(_y >= _rows) _y = _rows-1;				
		int _x = int(po1[i].x); if(_x < 0 ) _x = 0; if(_x >= _cols) _x = _cols-1;

		int id = label_map1.at<int>(_y, _x );

		Point2f _vec = po2[i] - po1[i];
		
		double len = sqrt(_vec.y * _vec.y  + _vec.x * _vec.x );

		if(len > frame_speed_limit * abs( frame_diff ))
			continue;

		double theta = atan2( _vec.y, _vec.x );

		mean_opt1[id] = mean_opt1[id] + Point2f( float(theta * len) , float(len) );

		count_kp[id] += 1.0f;
	}//end for matchcer

	for(int i = 0 ; i < n1 ; i++)
	{
		float sum_len = (mean_opt1[i].y + 1e-3f);				

		mean_opt1[i] = mean_opt1[i] * ( 1.0 / sum_len);

		sum_len/=count_kp[i];

		mean_opt1[i] = Point2f(  cos(mean_opt1[i].x) * sum_len,
								 sin(mean_opt1[i].x) * sum_len );
	}
}

void TimeSpaceMRF::getFlowMap(	vector< Point2i> & _maps,
								vector< Point2f > & mean_opt1,
								Mat & label_map1, Mat & label_map2,
								int n1,int n2)
{
	//========= affine into and set up link  ====================

	int _rows = label_map1.rows;
	int _cols = label_map1.cols;

	Mat _adj_mat = Mat::zeros( n1, n2 , CV_8U);

	int _count_link = 0;

	for(int y = 0 ; y < _rows; y++)
		for(int x = 0; x< _cols ; x++)
		{
			int id1 = label_map1.at< int >( y, x);

			int _y = y + int(mean_opt1[id1].y); if(_y < 0 ) _y = 0; if(_y >= _rows) _y = _rows-1;			
			int _x = x + int(mean_opt1[id1].x); if(_x < 0 ) _x = 0; if(_x >= _cols) _x = _cols-1;

			int id2 = label_map2.at< int >( _y , _x );

			if(!_adj_mat.at<uchar>(id1, id2))
			{
				_adj_mat.at<uchar>(id1, id2) = 1;
				_count_link++;
			}
		}

	_maps.reserve( _count_link);

	_maps.clear();

	for(int x = 0; x < n1; x++)
		for(int y = 0; y < n2 ; y++)
			if(_adj_mat.at<uchar>(x,y))
				_maps.push_back( Point2i( x, y) );
}



void TimeSpaceMRF::getSuperpixelFlow(	vector< SuperpixelFlow> & superpixel_flows,
										vector< Mat > & frms,
										vector< Mat >  & label_maps,
										vector< Mat >  & node_descs,
										vector< vector<Point2f> > & node_centers)
{//========= get optical flow  ====================

	vector<Match> matches;

	getMatches( frms, matches);

	superpixel_flows.resize( matches.size() );

	for(int match_id = 0; match_id < (int) matches.size() ; match_id++)
	{

		superpixel_flows[ match_id ].first = matches[match_id].first;

		int frame_id1 = matches[match_id].first.x;
		int frame_id2 = matches[match_id].first.y;

		vector< Point2f > & po1 = matches[match_id].second.first;
		vector< Point2f > & po2 = matches[match_id].second.second;

		Mat label_map1 = label_maps[ frame_id1];
		Mat label_map2 = label_maps[ frame_id2];

		Mat node_desc1 = node_descs[ frame_id1];
		Mat node_desc2 = node_descs[ frame_id2];			

		int n1 = node_desc1.rows;
		int n2 = node_desc2.rows;

		vector< Point2f > mean_opt1(n1, Point2f(0,0));

		getMeanOpt( mean_opt1, label_map1, po1, po2, n1, float(frame_id1 - frame_id2));

		vector< Point2i> & _maps = superpixel_flows[ match_id ].second;

		getFlowMap( _maps, mean_opt1, label_map1, label_map2,n1,n2);

		if(IF_VISUAL)
			saveOptFlow( frame_id1, frame_id2, node_centers[ frame_id1 ],mean_opt1 );

	}//end for each match		

}

void TimeSpaceMRF::saveOptFlow(	int frame_id1, int frame_id2,
								vector< Point2f > & _po,
								vector< Point2f > & mean_opt1)
{
	Mat dsp;
	char fname[400];
	sprintf( fname ,"./visual/SP_%d.jpg", frame_id1);
	dsp = imread( fname );

	if(_po.size() !=  mean_opt1.size())
	{
		cout << "warning! size of node unmatch with mean optical flow" << endl;
	}

	for(int i = 0 ; i < _po.size() ; i++)
	{
		cv::line( dsp , _po[i] , _po[i] + mean_opt1[i], CV_RGB( 255,0,0), 2,CV_AA);
	}

	sprintf( fname ,"./visual/track_%d_%d.jpg", frame_id1,frame_id2);

	imwrite( fname, dsp);
}



void TimeSpaceMRF::getGoodPoints(	vector< Mat > & frms,
									vector< vector< Point2f > > & key_points,
									LcKLT & my_tracker)
{
	key_points.resize( n_frame );

	int n = 1500;

	//========= get good keypoints  ====================

	for(int i = 0; i < n_frame; i++)
	{
		Mat _gray;
		cv::cvtColor( frms[i], _gray, CV_RGB2GRAY);
		my_tracker.getGoodPoints( _gray , key_points[i], n);
	}
}


void TimeSpaceMRF::getSLIC(	vector< Mat > & frms,
							vector< Mat > & likeli,
							vector< Mat >  & label_maps,
							vector< Mat >  & node_descs,
							vector< Mat >  & adj_mats,
							vector< int > & n_nodes,
							vector< int > & node_shift,
							vector< vector<Point2f> > & node_centers)
{
	label_maps.resize(n_frame);
	node_descs.resize(n_frame);
	adj_mats.resize(n_frame);
	n_nodes.resize(n_frame);
	node_shift.resize( n_frame );

	if(IF_VISUAL) node_centers.resize( n_frame );

	int node_count = 0;

	for(int i= 0 ; i < n_frame; i++)
	{
		LcSuperpixelCollector my_graphcut;
		my_graphcut.set( frms[i], likeli[i]);

		my_graphcut.collectNodeStatistic( frms[i], likeli[i] );
		my_graphcut.collectHeatmap( likeli[i] );

		label_maps[i] = my_graphcut.label_map;

		adj_mats[i] = my_graphcut.adj_mat;

		my_graphcut.computeAllNode( node_descs[i] );

		if(IF_VISUAL)
		{
			Mat dsp = my_graphcut.showSLIC( frms[i],label_maps[i] );

			char fname[400];
			sprintf( fname ,"./visual/SP_%d.jpg", i);
			imwrite( fname , dsp );				
		}

		node_shift[i] = node_count;

		n_nodes[i] = node_descs[i].rows;

		if(IF_VISUAL) node_centers[i] = my_graphcut.getNodeCenter( label_maps[i], n_nodes[i] );

		node_count += n_nodes[i];
	}

	cout << node_count << " superpixel in sequence " << endl;
}

void TimeSpaceMRF::getLikeli(	vector< Mat > & frms,
								vector< Mat > & likeli,
								LocalSensitiveBagging & crank_bagging,
								LcHistogram & global_extractor)
{	
	likeli.resize( n_frame );

	Mat empty_gt = Mat::Mat();

	for(int i = 0 ; i < n_frame ; i++)
		crank_bagging.predictFrame( frms[i], likeli[i], empty_gt, global_extractor);
}

void TimeSpaceMRF::saveSrc( vector< Mat > & frms )
{
	for(int i = 0 ; i < n_frame ; i++)
	{
		char fname[400];
		sprintf( fname ,"./visual/src_%d.jpg", i);
		imwrite( fname , frms[i] );
	}
}

void TimeSpaceMRF::saveLikeli( vector< Mat > & likeli )
{
	for(int i = 0 ; i < n_frame ; i++)
	{
		char fname[400];
		sprintf( fname ,"./visual/likelihood_%d.jpg", i);
		imwrite( fname , likeli[i] );
	}
}